#pragma once

#include "Object.hpp"
#include "Ray.hpp"
#include "Bounds3.hpp"
#include "Intersection.hpp"
#include "host_utils.hpp"
#include <algorithm>
#include <cassert>
#include <vector>
#include "UnifiedMemManaged.hpp"

struct BVHBuildNode : public UnifiedMemManaged
{
    Bounds3 bounds;
    BVHBuildNode *left;
    BVHBuildNode *right;
    Object *object;
    float area;

public:
    int splitAxis = 0, firstPrimOffset = 0, nPrimitives = 0;
    // BVHBuildNode Public Methods
    BVHBuildNode()
    {
        bounds = Bounds3(true);
        left = nullptr;
        right = nullptr;
        object = nullptr;
    }
};

// BVHAccel Forward Declarations
// struct BVHPrimitiveInfo;

// BVHAccel Declarations
static int leafNodes, totalLeafNodes, totalPrimitives, interiorNodes;
class BVHAccel : public UnifiedMemManaged
{
#define PRIMITIVES_MAX_NUM 24
public:
    // BVHAccel Public Types
    enum class SplitMethod
    {
        NAIVE,
        SAH
    };

    // BVHAccel Public Methods
    BVHAccel(Object *const p[], const int p_num, int maxPrimsInNode = 1, SplitMethod splitMethod = SplitMethod::NAIVE)
        : maxPrimsInNode(std::min(255, maxPrimsInNode)), splitMethod(splitMethod)
    {
        assert(primitives_num + p_num < PRIMITIVES_MAX_NUM);
        for (int i = 0; i < p_num; i++)
        {
            primitives[primitives_num++] = p[i];
        }

        std::vector<Object *> primitives_vector;
        for (int i = 0; i < primitives_num; i++)
        {
            primitives_vector.push_back(primitives[i]);
        }

        if (primitives_vector.empty())
            return;

        root = recursiveBuild(primitives_vector);
    }

    Bounds3 WorldBound() const;
    ~BVHAccel();

    __device__ Intersection Intersect(const Ray &ray) const;
    __device__ Intersection getIntersection(BVHBuildNode *node, const Ray &ray) const;
    //__device__ bool IntersectP(const Ray &ray) const;

    // BVHAccel Private Methods
    BVHBuildNode *recursiveBuild(std::vector<Object *> objects);

    // BVHAccel Private Data
    const int maxPrimsInNode;
    const SplitMethod splitMethod;
    Object *primitives[PRIMITIVES_MAX_NUM];
    int primitives_num = 0;
    BVHBuildNode *root;

    __device__ void getSample(BVHBuildNode *node, float p, Intersection &pos, float &pdf, int thread_id);
    __device__ void Sample(Intersection &pos, float &pdf, int thread_id);
};
